#!/usr/bin/env python3

import asyncio
import sys
import os
import struct
import base64
import hashlib


class WebSocketClient:
    def __init__(self, host, port, path):
        self.host = host
        self.port = port
        self.path = path
        self.reader = None
        self.writer = None
        self.is_connected = False

    async def connect(self):
        """建立WebSocket连接"""
        try:
            # 建立TCP连接
            self.reader, self.writer = await asyncio.open_connection(self.host, self.port)
            
            # 生成WebSocket握手密钥
            key = base64.b64encode(os.urandom(16)).decode()
            
            # 发送WebSocket握手请求
            handshake = f"""GET {self.path} HTTP/1.1\r
Host: {self.host}:{self.port}\r
Upgrade: websocket\r
Connection: Upgrade\r
Sec-WebSocket-Key: {key}\r
Sec-WebSocket-Version: 13\r
\r
"""
            self.writer.write(handshake.encode())
            await self.writer.drain()
            
            # 读取服务器响应
            response = await self.reader.read(1024)
            response_str = response.decode()
            
            # 验证是否握手成功
            if "101 Switching Protocols" in response_str and "Upgrade: websocket" in response_str:
                print(f"已连接到 ws://{self.host}:{self.port}{self.path}")
                self.is_connected = True
                # 调用on_open回调
                if hasattr(self, 'on_open') and callable(self.on_open):
                    self.on_open()
                return True
            else:
                print(f"WebSocket握手失败: {response_str}")
                return False
        except Exception as e:
            print(f"连接错误: {e}")
            return False

    def create_frame(self, opcode, data):
        """创建WebSocket帧"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        
        length = len(data)
        frame = bytearray()
        
        # FIN位(1) + RSV1-3(0) + opcode
        frame.append(0x80 | opcode)
        
        # Mask位(1) + 数据长度
        if length < 126:
            frame.append(0x80 | length)
        elif length < 65536:
            frame.append(0x80 | 126)
            frame.extend(struct.pack('!H', length))
        else:
            frame.append(0x80 | 127)
            frame.extend(struct.pack('!Q', length))
        
        # Mask key (客户端必须掩码)
        mask_key = os.urandom(4)
        frame.extend(mask_key)
        
        # 掩码数据
        masked_data = bytearray()
        for i in range(length):
            masked_data.append(data[i] ^ mask_key[i % 4])
        
        frame.extend(masked_data)
        return frame

    async def send_text(self, message):
        """发送文本消息"""
        if self.writer and self.is_connected:
            frame = self.create_frame(0x1, message)
            self.writer.write(frame)
            await self.writer.drain()

    async def send_binary(self, data):
        """发送二进制消息"""
        if self.writer and self.is_connected:
            frame = self.create_frame(0x2, data)
            self.writer.write(frame)
            await self.writer.drain()

    def parse_frame(self, data):
        """解析WebSocket帧"""
        if len(data) < 2:
            return None, None, None, None
            
        # 解析第一个字节
        fin = (data[0] & 0x80) != 0
        opcode = data[0] & 0x0F
        
        # 解析第二个字节
        masked = (data[1] & 0x80) != 0
        payload_length = data[1] & 0x7F
        
        offset = 2
        
        # 解析长度
        if payload_length == 126:
            if len(data) < offset + 2:
                return None, None, None, None
            payload_length = struct.unpack('!H', data[offset:offset+2])[0]
            offset += 2
        elif payload_length == 127:
            if len(data) < offset + 8:
                return None, None, None, None
            payload_length = struct.unpack('!Q', data[offset:offset+8])[0]
            offset += 8
        
        # 解析掩码键
        if masked:
            if len(data) < offset + 4:
                return None, None, None, None
            mask_key = data[offset:offset+4]
            offset += 4
        
        # 解析数据
        if len(data) < offset + payload_length:
            return None, None, None, None
            
        payload = data[offset:offset+payload_length]
        
        # 如果有掩码则解掩码
        if masked:
            for i in range(len(payload)):
                payload[i] ^= mask_key[i % 4]
        
        return fin, opcode, payload, offset + payload_length

    async def receive_messages(self):
        """接收并处理消息"""
        buffer = bytearray()
        while self.is_connected:
            try:
                data = await self.reader.read(4096)
                if not data:
                    self.is_connected = False
                    print("[-] tcp connection closed")
                    self.on_close()
                    break
                    
                buffer.extend(data)
                
                # 解析帧
                while len(buffer) > 0:
                    fin, opcode, payload, consumed = self.parse_frame(buffer)
                    if fin is None:
                        # 数据不足，等待更多数据
                        break
                        
                    # 移除已解析的数据
                    buffer = buffer[consumed:]
                    
                    # 处理不同类型的操作码
                    if opcode == 0x1:  # 文本帧
                        message = payload.decode('utf-8')
                        # 调用on_msg回调
                        if hasattr(self, 'on_msg') and callable(self.on_msg):
                            self.on_msg(message, is_binary=False)
                    elif opcode == 0x2:  # 二进制帧
                        # 调用on_msg回调
                        if hasattr(self, 'on_msg') and callable(self.on_msg):
                            self.on_msg(payload, is_binary=True)
                    elif opcode == 0x8:  # 关闭帧
                        self.is_connected = False
                        print("[-] WebSocket connection closed")
                        # 调用on_close回调
                        if hasattr(self, 'on_close') and callable(self.on_close):
                            self.on_close()
                        return
                    elif opcode == 0x9:  # Ping帧
                        # 发送Pong响应
                        pong_frame = self.create_frame(0xA, payload)
                        self.writer.write(pong_frame)
                        await self.writer.drain()
                        
            except Exception as e:
                print(f"接收消息时出错: {e}")
                self.is_connected = False
                break

    async def close(self):
        """关闭连接"""
        self.is_connected = False
        if self.writer:
            # 发送关闭帧
            close_frame = self.create_frame(0x8, b'')
            try:
                self.writer.write(close_frame)
                await self.writer.drain()
            except:
                pass
            self.writer.close()
            await self.writer.wait_closed()


class ShellClient:
    def __init__(self, ws_client: WebSocketClient):
        self.ws_client = ws_client
        self.ws_client.on_open = self.on_open
        self.ws_client.on_msg = self.on_message
        self.ws_client.on_close = self.on_close

    def on_open(self):
        """WebSocket连接打开时调用"""
        print("现在可以输入命令，输入 'exit' 退出")
        print("-" * 40)

    def on_message(self, message, is_binary=False):
        """收到消息时调用"""
        if is_binary:
            # 二进制消息直接写入标准输出
            sys.stdout.buffer.write(message)
            sys.stdout.flush()
        else:
            # 文本消息
            print(f"{message}", end='', flush=True)

    def on_close(self):
        """连接关闭时调用"""
        print("connection closed")
        self.input_task.cancel()

    async def stdin_loop(self):
        """发送用户输入"""
        loop = asyncio.get_event_loop()
        while self.ws_client.is_connected:
            try:
                # 使用线程来避免阻塞事件循环
                user_input = await loop.run_in_executor(None, sys.stdin.readline)
                if user_input and self.ws_client.is_connected:
                    await self.ws_client.send_text(user_input)
                else:
                    break
            except asyncio.exceptions.CancelledError:
                print(f"stdin loop cancled") 
                break

    async def run(self):
        """运行客户端"""
        self.input_task = asyncio.create_task(self.stdin_loop())
        self.ws_message_task = asyncio.create_task(self.ws_client.receive_messages())
        try:
            await asyncio.gather(self.input_task, self.ws_message_task)
            print("all task done")
        except KeyboardInterrupt:
            print("\n客户端已关闭")
        finally:
            if self.ws_client.is_connected:  # 如果连接已关闭，则取消输入任务
                await self.ws_client.close()
            print('bye')


async def main():
    if len(sys.argv) != 2:
        print("用法: python ws_client.py <服务器地址:端口>")
        print("例如: python ws_client.py localhost:8080")
        return
    
    # 创建WebSocket客户端
    ip_port = sys.argv[1]
    ws_client = WebSocketClient('localhost', 8080, f'/ws/shell?target={ip_port}')
    # 创建Shell客户端
    shell_client = ShellClient(ws_client)
    # 连接
    if not await ws_client.connect():
        print("无法连接到服务器")
        return
    print(f"connected to websocket ")
    
    # 运行
    await shell_client.run()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n👋 服务器已停止")